# 使用训练得到的模型进行椭圆拟合
import argparse
import os
import cv2
import torch
import numpy as np
import json

from angle_calculation import calculate_antervision, calculate_p, fit_ellipse_and_get_major_axis, calculate_angle
from UnetDatasets import UnetDataset
from UnetModel import UnetModel
from config import dataset_path, train_config, project_path


# 预测
def evaluate(epoch, split):
    if epoch == 0:
        epoch = "min_loss"
    test_dataset = UnetDataset(dataset_path=dataset_path, split=split)

    # 定义模型
    model = UnetModel()
    model.to(train_config['device'])

    save_dir = os.path.join(project_path, "result", "weights")
    img_save_dir = os.path.join(project_path, "result", "ellipse_img", split)
    if not os.path.exists(img_save_dir):
        os.makedirs(img_save_dir)

    model_weight_path = os.path.join(save_dir, str(epoch), f"epoch_{str(epoch)}.pkl")
    model.load_state_dict(torch.load(model_weight_path))
    model.eval()

    for index in range(len(test_dataset)):
        img, target_point, img_base, img_name_and_path, window_size = test_dataset.getToVisualization(index=index)
        print(f"index:{index}, file:{img_name_and_path}")
        img = torch.unsqueeze(img, dim=0)  # 训练时采用的是DataLoader函数, 会直接增加第一个维度
        img = img.to(train_config['device'])
        pre = model(img)
        pre = pre.cpu().detach().numpy()
        pre = pre.reshape(pre.shape[0], -1, 2)

        target_point = torch.unsqueeze(target_point, dim=0)
        target_point = target_point.reshape(target_point.shape[0], -1, 2)
        # print(pre.shape)
        # print(target_point.shape)
        # mask掉没有填充骨头的
        mask = (target_point == 0)
        pre[mask] = 0
        # 获取第一张图片的名称
        img_name = os.path.split(img_name_and_path)[1]
        # 计算并保存前倾角
        save_antervision(pre[0], img_name.replace('.bmp', '.txt'), window_size)
        # 标注椭圆
        img_with_ellipse = show_ellipse_on_picture(img_base, pre[0], target_point[0], window_size)  # 传入第一个图像的预测结果
        # 显示或保存图像
        cv2.imwrite(os.path.join(img_save_dir, img_name.replace('.bmp', '.jpg')), img_with_ellipse)
        # cv2.imshow('Image with Keypoints', img_with_ellipse)
        # cv2.waitKey(0)  # 按任意键关闭窗口
        # cv2.destroyAllWindows()


def show_ellipse_on_picture(img, landmarks, landmarks_gt=None, window_size=None):
    (width, height) = window_size
    # 提取右侧和左侧的关键点
    right_points, left_points = extract_points(landmarks, window_size)

    # 对右侧和左侧关键点进行椭圆拟合并绘制
    img = fit_and_draw_ellipse(right_points, img, (0, 255, 0))  # 绿色椭圆表示右侧
    img = fit_and_draw_ellipse(left_points, img, (255, 0, 0))  # 红色椭圆表示左侧

    # 绘制预测的关键点
    for point in landmarks:
        point = tuple([int(point[0] / 512 * width), int(point[1] / 512 * height)])
        # 对于不好的模型，radius = 12太大了，会遮住椭圆
        img = cv2.circle(img, center=point, radius=12, color=(0, 0, 255), thickness=-1)

        # 绘制真实标签，直观显示拟合程度
    if landmarks_gt is not None:
        for point in landmarks_gt:
            point = tuple([int(point[0] / 512 * width), int(point[1] / 512 * height)])
            img = cv2.circle(img, center=point, radius=12, color=(0, 255, 0), thickness=-1)

    return img


def extract_points(landmarks, window_size):
    (width, height) = window_size
    right_points = []
    left_points = []
    for index in range(0, len(landmarks)):
        point = landmarks[index]
        point = tuple([int(point[0] / 512 * width), int(point[1] / 512 * height)])
        # 没填充的点不显示
        if point == (0, 0):
            continue
        # 右侧关键点为1-6
        if index in range(0, 6):
            right_points.append(point)
        # 左侧关键点为8-13
        elif index in range(7, 13):
            left_points.append(point)
    return np.array(right_points, dtype=np.float32), np.array(left_points, dtype=np.float32)


# 拟合椭圆并绘制
def fit_and_draw_ellipse(points, img, color):
    if len(points) >= 5:  # 椭圆拟合需要至少 5 个点
        ellipse = cv2.fitEllipse(points)
        cv2.ellipse(img, ellipse, color, 2)
    # else:
        # print("Not enough points to fit an ellipse.")
    return img

def save_antervision(landmarks, save_name, window_size):
    # 提取关键点 1~6 作为右侧点，8~13 作为左侧点
    right_points, left_points = extract_points(landmarks, window_size)
    # if len(right_points) < 2 or len(left_points) < 2:
    #     print(f"Skipping {save_path} due to insufficient points for abduction angle calculation.")
    #     return

    # 计算前倾角（拟合椭圆并使用前倾角公式）
    right_major_axis, right_center, right_ellipse = fit_ellipse_and_get_major_axis(right_points, save_name, "右侧")
    left_major_axis, left_center, left_ellipse = fit_ellipse_and_get_major_axis(left_points, save_name, "左侧")

    antervision_right = calculate_antervision(right_major_axis, right_ellipse) if right_major_axis is not None else None
    antervision_left = calculate_antervision(left_major_axis, left_ellipse) if left_major_axis is not None else None

    # 计算右侧髋关节的外展角（检查是否有足够的点）
    if len(right_points) >= 6:
        abduction_angle_right = calculate_angle(right_points[1] - right_points[0], right_points[5] - right_points[4])  # 使用1~2 和5~6
    elif len(right_points) >= 2:  # 如果不足6个点，用最靠近的两个点
        abduction_angle_right = calculate_angle(right_points[1] - right_points[0], right_points[-1] - right_points[-2])
    else:
        abduction_angle_right = None

    # 计算左侧髋关节的外展角（检查是否有足够的点）
    if len(left_points) >= 6:
        abduction_angle_left = calculate_angle(left_points[1] - left_points[0], left_points[5] - left_points[4])  # 使用8~9 和12~13
    elif len(left_points) >= 2:  # 如果不足6个点，用最靠近的两个点
        abduction_angle_left = calculate_angle(left_points[1] - left_points[0], left_points[-1] - left_points[-2])
    else:
        abduction_angle_left = None

    # 保存计算结果到文件
    save_angles_to_file(save_name, abduction_angle_right, abduction_angle_left, antervision_right, antervision_left)


# 保存角度信息到txt文件（中文输出）
def save_angles_to_file(save_name, abduction_angle_right, abduction_angle_left, antervision_right, antervision_left):
    # 构建结果保存路径
    output_folder = "./result/module_angle"
    os.makedirs(output_folder, exist_ok=True)

    # 创建文件名
    filename = os.path.basename(save_name).replace('.json', '.txt')
    output_path = os.path.join(output_folder, filename)

    # 写入文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(f"文件：{save_name}\n")

        # 外展角
        if abduction_angle_right is not None:
            f.write(f"右髋关节外展角：{abduction_angle_right:.4f} 弧度\n")
        else:
            f.write("右髋关节外展角：未计算\n")

        if abduction_angle_left is not None:
            f.write(f"左髋关节外展角：{abduction_angle_left:.4f} 弧度\n")
        else:
            f.write("左髋关节外展角：未计算\n")

        # 前倾角
        if antervision_right is not None:
            f.write(f"右髋关节前倾角：{antervision_right:.4f} 弧度\n")
        else:
            f.write("右髋关节前倾角：未计算\n")

        if antervision_left is not None:
            f.write(f"左髋关节前倾角：{antervision_left:.4f} 弧度\n")
        else:
            f.write("左髋关节前倾角：未计算\n")

    print(f"角度信息已保存到 {output_path}")


def run():
    parser = argparse.ArgumentParser()
    parser.add_argument('--epoch', type=int, help='输出调用指定的轮次，默认调用min_loss的', default=0)
    parser.add_argument('--split', type=str, help='使用train还是test进行测试', default="test", choices=["test", "train"])
    # 对一组权重进行预测
    # 解析命令行参数
    args = parser.parse_args()
    print(args)
    evaluate(epoch=args.epoch, split=args.split)


if __name__ == "__main__":
    run()
